package wiki.xsx.core.bytecode.transformer;

import javassist.CtBehavior;
import javassist.CtClass;
import javassist.bytecode.ConstPool;
import javassist.bytecode.MethodInfo;
import javassist.bytecode.ParameterAnnotationsAttribute;
import javassist.bytecode.annotation.Annotation;
import javassist.bytecode.annotation.MemberValue;
import lombok.SneakyThrows;
import wiki.xsx.core.bytecode.enums.ExecuteType;
import wiki.xsx.core.bytecode.model.Annotations;
import wiki.xsx.core.bytecode.model.ParameterAnnotations;

import java.util.*;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 抽象方法参数注解转换器
 *
 * @author xsx
 * @date 2023/3/2
 * @since 1.8
 * <p>
 * Copyright (c) 2023 xsx All Rights Reserved.
 * easy-bytecode is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 * http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 * </p>
 */
public abstract class AbstractMethodParamAnnotationTransformer extends AbstractAnnotationTransformer {

    /**
     * 方法参数类型
     */
    protected List<Class<?>> parameterTypes;
    /**
     * 执行类型
     */
    protected ExecuteType executeType;
    /**
     * 注解数组
     */
    protected ParameterAnnotations[] annotations;

    /**
     * 有参构造
     *
     * @param isWrite        是否写入
     * @param fileDir        文件目录
     * @param className      类名称
     * @param executeType    执行类型
     * @param methodName     方法名称
     * @param parameterTypes 参数类型
     * @param annotations    注解
     */
    public AbstractMethodParamAnnotationTransformer(
            Boolean isWrite,
            String fileDir,
            ExecuteType executeType,
            String className,
            String methodName,
            List<Class<?>> parameterTypes,
            ParameterAnnotations... annotations
    ) {
        super(isWrite, fileDir, className, methodName);
        this.parameterTypes = Optional.ofNullable(parameterTypes).orElse(new ArrayList<>(0));
        this.executeType = executeType;
        this.annotations = Optional.ofNullable(annotations).orElse(new ParameterAnnotations[0]);
    }

    /**
     * 处理
     *
     * @param ctClass    ct类
     * @param ctBehavior ct行为
     */
    protected void process(CtClass ctClass, CtBehavior ctBehavior) {
        this.process(ctClass, ctBehavior.getMethodInfo());
    }

    /**
     * 处理
     *
     * @param ctClass    ct类
     * @param methodInfo 方法信息
     */
    @SneakyThrows
    private void process(CtClass ctClass, MethodInfo methodInfo) {
        ConstPool constPool = ctClass.getClassFile().getConstPool();
        switch (this.executeType) {
            // 新增或更新
            case ANNOTATION_METHOD_PARAM_ADD_OR_UPDATE: {
                this.processAnnotation(methodInfo, constPool, this::addOrUpdateAnnotation);
                break;
            }
            // 移除
            case ANNOTATION_METHOD_PARAM_REMOVE: {
                this.processAnnotation(methodInfo, constPool, this::removeAnnotation);
                break;
            }
        }
    }

    /**
     * 初始化参数注解
     *
     * @return 返回参数注解字典
     */
    private Map<Annotations.Policy, List<ParameterAnnotations>> initParameterAnnotations() {
        Map<Annotations.Policy, List<ParameterAnnotations>> map = new HashMap<>(2);
        List<ParameterAnnotations> visibleList = new ArrayList<>(this.annotations.length);
        List<ParameterAnnotations> invisibleList = new ArrayList<>(this.annotations.length);
        for (ParameterAnnotations parameterAnnotations : this.annotations) {
            List<Annotations> newList = parameterAnnotations.getAnnotations();
            if (newList == null) {
                newList = new ArrayList<>(0);
            }
            List<Annotations> visibleAnnotationList = new ArrayList<>();
            List<Annotations> invisibleAnnotationlist = new ArrayList<>();
            for (Annotations annotation : newList) {
                if (annotation.getPolicy() == Annotations.Policy.RuntimeVisible) {
                    visibleAnnotationList.add(annotation);
                } else {
                    invisibleAnnotationlist.add(annotation);
                }
            }
            Integer index = parameterAnnotations.getIndex();
            if (!visibleAnnotationList.isEmpty()) {
                visibleList.add(new ParameterAnnotations(index, visibleAnnotationList));
            }
            if (!invisibleAnnotationlist.isEmpty()) {
                invisibleList.add(new ParameterAnnotations(index, invisibleAnnotationlist));
            }
        }
        map.put(Annotations.Policy.RuntimeVisible, visibleList);
        map.put(Annotations.Policy.RuntimeInvisible, invisibleList);
        return map;
    }

    /**
     * 获取属性
     *
     * @param info      方法信息
     * @param constPool 常量池
     * @param policy    策略
     * @return 返回属性
     */
    private ParameterAnnotationsAttribute getAttribute(MethodInfo info, ConstPool constPool, Annotations.Policy policy) {
        String tag = policy == Annotations.Policy.RuntimeVisible ? ParameterAnnotationsAttribute.visibleTag : ParameterAnnotationsAttribute.invisibleTag;
        ParameterAnnotationsAttribute attribute = (ParameterAnnotationsAttribute) info.getAttribute(tag);
        if (attribute == null) {
            attribute = new ParameterAnnotationsAttribute(constPool, tag);
        }
        return attribute;
    }

    /**
     * 处理注解
     *
     * @param methodInfo 方法信息
     * @param constPool  常量池
     * @param function   处理方法
     */
    private void processAnnotation(MethodInfo methodInfo, ConstPool constPool, ParameterAnnotationFunction<ConstPool, ParameterAnnotationsAttribute, ParameterAnnotations> function) {
        Map<Annotations.Policy, List<ParameterAnnotations>> map = this.initParameterAnnotations();
        AtomicReference<ParameterAnnotationsAttribute> attribute = new AtomicReference<>();
        map.forEach(
                (policy, value) -> {
                    value.forEach(
                            annotation -> {
                                attribute.set(this.getAttribute(methodInfo, constPool, policy));
                                function.process(constPool, attribute.get(), annotation);
                            }
                    );
                    Optional.ofNullable(attribute.get()).ifPresent(
                            attr -> {
                                methodInfo.addAttribute(attr);
                                attribute.set(null);
                            }
                    );
                }
        );
    }

    /**
     * 添加注解
     *
     * @param constPool   常量池
     * @param attribute   参数属性
     * @param annotations 参数注解
     */
    @SneakyThrows
    private void addOrUpdateAnnotation(
            ConstPool constPool,
            ParameterAnnotationsAttribute attribute,
            ParameterAnnotations annotations
    ) {
        Annotation[][] oldAnnotations = attribute.getAnnotations();
        if (oldAnnotations.length == 0) {
            oldAnnotations = Arrays.stream(new Annotation[this.parameterTypes.size()][]).map(v -> new Annotation[0]).toArray(Annotation[][]::new);
        }
        List<Annotations> newList = annotations.getAnnotations();
        if (newList == null) {
            return;
        }
        Integer index = annotations.getIndex();
        Annotation[] oldArray = oldAnnotations[index];
        List<Annotation> list;
        if (oldArray == null || oldArray.length == 0) {
            list = new ArrayList<>(newList.size());
            for (Annotations value : newList) {
                list.add(this.createAnnotation(constPool, value));
            }
        } else {
            Map<String, Annotation> oldMap = Arrays.stream(oldArray).collect(Collectors.toMap(Annotation::getTypeName, Function.identity()));
            int count = newList.size() + oldArray.length;
            List<Annotation> addList = new ArrayList<>(count);
            list = new ArrayList<>(count);
            for (Annotations newAnnotation : newList) {
                Annotation oldAnnotation = oldMap.get(newAnnotation.getName());
                if (oldAnnotation == null) {
                    addList.add(this.createAnnotation(constPool, newAnnotation));
                } else {
                    this.updateAnnotation(constPool, oldAnnotation, newAnnotation);
                }
            }
            list.addAll(Arrays.asList(oldArray));
            list.addAll(addList);
        }
        oldAnnotations[index] = list.toArray(new Annotation[0]);
        attribute.setAnnotations(oldAnnotations);
    }

    /**
     * 创建注解
     *
     * @param constPool   常量池
     * @param annotations 注解
     * @return 返回新注解
     */
    private Annotation createAnnotation(ConstPool constPool, Annotations annotations) {
        Annotation annotation = new Annotation(annotations.getName(), constPool);
        Optional.ofNullable(annotations.getAttributes()).ifPresent(
                attrs -> attrs.forEach(
                        attr -> annotation.addMemberValue(
                                attr.getName(), this.createMemberValue(constPool, attr)
                        )
                )
        );
        return annotation;
    }

    /**
     * 更新注解
     *
     * @param constPool     常量池
     * @param oldAnnotation 原注解
     * @param newAnnotation 新注解
     */
    private void updateAnnotation(ConstPool constPool, Annotation oldAnnotation, Annotations newAnnotation) {
        if (oldAnnotation == null) {
            return;
        }
        Set<String> names = oldAnnotation.getMemberNames();
        Map<String, Annotations.Attribute> attributes = newAnnotation.getAttributes().stream().collect(Collectors.toMap(Annotations.Attribute::getName, Function.identity()));
        for (String name : names) {
            MemberValue memberValue = oldAnnotation.getMemberValue(name);
            Annotations.Attribute attr = attributes.get(name);
            if (attr != null) {
                memberValue = this.createMemberValue(constPool, attr);
            }
            oldAnnotation.addMemberValue(name, memberValue);
        }
    }

    /**
     * 移除注解
     *
     * @param constPool   常量池
     * @param attribute   属性
     * @param annotations 注解
     */
    private void removeAnnotation(
            ConstPool constPool,
            ParameterAnnotationsAttribute attribute,
            ParameterAnnotations annotations
    ) {
        Annotation[][] oldAnnotations = attribute.getAnnotations();
        if (oldAnnotations.length == 0) {
            return;
        }
        Integer index = annotations.getIndex();
        Annotation[] oldArray = oldAnnotations[index];
        if (oldArray == null || oldArray.length == 0) {
            return;
        } else {
            Map<String, Annotations> newMap = annotations.getAnnotations().stream().collect(Collectors.toMap(Annotations::getName, Function.identity()));
            List<Annotation> list = new ArrayList<>(oldArray.length);
            for (Annotation oldAnnotation : oldArray) {
                if (newMap.get(oldAnnotation.getTypeName()) == null) {
                    list.add(oldAnnotation);
                }
            }
            oldAnnotations[index] = list.toArray(new Annotation[0]);
        }
        attribute.setAnnotations(oldAnnotations);
    }

    /**
     * 注解函数式接口
     *
     * @param <ConstPool>                     常量池
     * @param <ParameterAnnotationsAttribute> 属性
     * @param <ParameterAnnotations>          注解
     */
    @FunctionalInterface
    private interface ParameterAnnotationFunction<ConstPool, ParameterAnnotationsAttribute, ParameterAnnotations> {
        /**
         * 处理
         *
         * @param constPool   常量池
         * @param attribute   属性
         * @param annotations 注解
         */
        void process(ConstPool constPool, ParameterAnnotationsAttribute attribute, ParameterAnnotations annotations);
    }
}
